#include <iostream>

#include "net.hpp"

int main(int argc, char* argv[])
{
    const int numSegement = 3;
    const int numPoints   = 201;
    std::vector<float> rangeX{0, 4};
    const float dx  = (rangeX[1] - rangeX[0]) / (numPoints - 1);
    torch::Tensor x = torch::zeros({numPoints, 1});
    torch::Tensor y = torch::zeros({numPoints, 1});
    for (int i = 0; i < numPoints; ++i)
    {
        const float x_ = rangeX[0] + i * dx;
        x[i][0]        = x_;
        if (x_ < 1.0)
        {
            y[i][0] = x_;
        }
        else if (x_ > 3.0)
        {
            y[i][0] = x_ - 4.0;
        }
        else
        {
            y[i][0] = 2.0 - x_;
        }
        const auto random = torch::randn({1}) * 0.1;
        y[i][0] += random.item<float>();
    }

    std::shared_ptr<Net> net = std::make_shared<Net>(numSegement, rangeX);

    auto optimizer = std::make_shared<torch::optim::AdamW>(net->parameters(), 1.e-3);

    auto dataSet    = CustomDtaset(x, y).map(torch::data::transforms::Stack<>());
    auto dataLoader = torch::data::make_data_loader<torch::data::samplers::RandomSampler>(
        std::move(dataSet), 32);

    float lossVal = 1.0, oldLoss = 1.0;
    const int maxEpoches = 2000;
    int epoch            = 0;
    std::cout << std::scientific;
    torch::Tensor zerosContinuity = torch::zeros({numSegement - 1, 1});
    while (std::abs(lossVal) > 1.e-6)
    {
        lossVal = 0.0;
        net->train();
        for (auto& batch : *dataLoader)
        {
            torch::Tensor outD   = net->forward(batch.data);
            torch::Tensor target = torch::cat({batch.target, zerosContinuity}, 0);
            torch::Tensor lossD  = torch::mse_loss(outD, target);
            optimizer->zero_grad();
            lossD.backward();
            optimizer->step();

            //
            lossVal += lossD.item<float>();
        }
        epoch++;

        net->checkSegementPoint();

        if (std::abs(lossVal / oldLoss - 1.0) < 5.e-6) break;
        std::cout << "Epoch.Idx : " << epoch << ", Train Loss : " << lossVal
                  << ", Segement Point: ";
        net->printSegementPoint();
        oldLoss = lossVal;

        net->eval();
        std::ofstream os("data_from_train_" + std::to_string(epoch) + ".csv");
        torch::Tensor out = net->forward(x);
        for (int i = 0; i < numPoints; ++i)
        {
            os << x[i][0].item<float>() << " " << y[i][0].item<float>() << " "
               << out[i][0].item<float>() << std::endl;
        }
        os.close();

        if (epoch >= maxEpoches) break;
    }

    net->printSegementPoint();
    net->eval();
    std::ofstream os("data_from_train.csv");
    torch::Tensor out = net->forward(x);
    for (int i = 0; i < numPoints; ++i)
    {
        os << x[i][0].item<float>() << "," << y[i][0].item<float>() << ","
           << out[i][0].item<float>() << std::endl;
    }
    os.close();

    return 0;
}
